#!/bin/bash
# Copyright 2018-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

############################################################################
## This script generates the model repository needed by some of the
## tritonserver CI tests. Generating these models requires using
## the TensorFlow and PyTorch containers.
##
## 1. Update TENSORRT_IMAGE, PYTORCH_IMAGE and TENSORFLOW_IMAGE to
## match what is being used by the tritonserver release being
## tested.
##
## 2. Set CUDA_DEVICE to the ID of the CUDA device present on the
## local system that you want to target for the generated models.
##
## 3. Run this script to create /tmp/qa_model_repository,
## /tmp/qa_variable_model_repository, /tmp/qa_shapetensor_model_repository
## /tmp/qa_identity_model_repository, /tmp/qa_identity_big_model_repository
## /tmp/qa_reshape_model_repository, /tmp/qa_noshape_model_repository,
## /tmp/qa_sequence_model_repository, /tmp/qa_ensemble_model_repository
## /tmp/qa_dyna_sequence_model_repository, and
## /tmp/qa_variable_sequence_model_repository directories containing
## all the models needed for CI testing.
##
############################################################################

TRITON_VERSION=${TRITON_VERSION:=24.07}

# ONNX. Use ONNX_OPSET 0 to use the default for ONNX version
ONNX_VERSION=1.13.0
ONNX_OPSET=0

# OPENVINO version
OPENVINO_VERSION=2023.3.0

UBUNTU_IMAGE=${UBUNTU_IMAGE:=ubuntu:22.04}
PYTORCH_IMAGE=${PYTORCH_IMAGE:=nvcr.io/nvidia/pytorch:$TRITON_VERSION-py3}
TENSORFLOW_IMAGE=${TENSORFLOW_IMAGE:=nvcr.io/nvidia/tensorflow:$TRITON_VERSION-tf2-py3}
TENSORRT_IMAGE=${TENSORRT_IMAGE:=nvcr.io/nvidia/tensorrt:$TRITON_VERSION-py3}
CUDA_DEVICE=${NV_GPU:=0}

[[ $RUNNER_GPUS =~ ^[0-9] ]] && DOCKER_GPU_ARGS=$(eval $NV_DOCKER_ARGS) || DOCKER_GPU_ARGS="--gpus device=$CUDA_DEVICE"

###
HOST_BUILD_DIR=${HOST_BUILD_DIR:=/tmp}
HOST_SRCDIR=$HOST_BUILD_DIR/gen_srcdir
HOST_DESTDIR=$HOST_BUILD_DIR/$TRITON_VERSION/qa_model_repository
HOST_VARDESTDIR=$HOST_BUILD_DIR/$TRITON_VERSION/qa_variable_model_repository
HOST_IDENTITYDESTDIR=$HOST_BUILD_DIR/$TRITON_VERSION/qa_identity_model_repository
HOST_SIGDEFDESTDIR=$HOST_BUILD_DIR/$TRITON_VERSION/qa_tf_tag_sigdef_repository
HOST_IDENTITYBIGDESTDIR=$HOST_BUILD_DIR/$TRITON_VERSION/qa_identity_big_model_repository
HOST_TFPARAMETERSDESTDIR=$HOST_BUILD_DIR/$TRITON_VERSION/qa_tf_parameters_repository
HOST_SHAPEDESTDIR=$HOST_BUILD_DIR/$TRITON_VERSION/qa_shapetensor_model_repository
HOST_RESHAPEDESTDIR=$HOST_BUILD_DIR/$TRITON_VERSION/qa_reshape_model_repository
HOST_SEQDESTDIR=$HOST_BUILD_DIR/$TRITON_VERSION/qa_sequence_model_repository
HOST_DYNASEQDESTDIR=$HOST_BUILD_DIR/$TRITON_VERSION/qa_dyna_sequence_model_repository
HOST_DYNASEQIMPLICITDESTDIR=$HOST_BUILD_DIR/$TRITON_VERSION/qa_dyna_sequence_implicit_model_repository
HOST_VARSEQDESTDIR=$HOST_BUILD_DIR/$TRITON_VERSION/qa_variable_sequence_model_repository
HOST_ENSEMBLEDESTDIR=$HOST_BUILD_DIR/$TRITON_VERSION/qa_ensemble_model_repository
HOST_NOSHAPEDESTDIR=$HOST_BUILD_DIR/$TRITON_VERSION/qa_noshape_model_repository
HOST_PLGDESTDIR=$HOST_BUILD_DIR/$TRITON_VERSION/qa_trt_plugin_model_repository
HOST_RAGGEDDESTDIR=$HOST_BUILD_DIR/$TRITON_VERSION/qa_ragged_model_repository
HOST_FORMATDESTDIR=$HOST_BUILD_DIR/$TRITON_VERSION/qa_trt_format_model_repository
HOST_DATADEPENDENTDIR=$HOST_BUILD_DIR/$TRITON_VERSION/qa_trt_data_dependent_model_repository
HOST_IMPLICITSEQDESTDIR=$HOST_BUILD_DIR/$TRITON_VERSION/qa_sequence_implicit_model_repository
HOST_VARIMPLICITSEQDESTDIR=$HOST_BUILD_DIR/$TRITON_VERSION/qa_variable_sequence_implicit_model_repository
HOST_INITIALSTATEIMPLICITSEQDESTDIR=$HOST_BUILD_DIR/$TRITON_VERSION/qa_sequence_initial_state_implicit_model_repository
HOST_VARINITIALSTATEIMPLICITSEQDESTDIR=$HOST_BUILD_DIR/$TRITON_VERSION/qa_variable_sequence_initial_state_implicit_model_repository
HOST_TORCHTRTDESTDIR=$HOST_BUILD_DIR/$TRITON_VERSION/torchtrt_model_store
HOST_SCALARMODELSDESTDIR=$HOST_BUILD_DIR/$TRITON_VERSION/qa_scalar_models
HOST_IMAGEMODELSDESTDIR=$HOST_BUILD_DIR/$TRITON_VERSION/qa_dynamic_batch_image_model_repository

rm -fr $HOST_SRCDIR $HOST_DESTDIR $HOST_VARDESTDIR $HOST_TFPARAMETERSDESTDIR
rm -fr $HOST_IDENTITYDESTDIR $HOST_IDENTITYBIGDESTDIR $HOST_SHAPEDESTDIR $HOST_SIGDEFDESTDIR
rm -fr $HOST_SEQDESTDIR $HOST_DYNASEQDESTDIR $HOST_VARSEQDESTDIR
rm -fr $HOST_ENSEMBLEDESTDIR $HOST_NOSHAPEDESTDIR $HOST_RESHAPEDESTDIR
rm -fr $HOST_PLGDESTDIR $HOST_RAGGEDDESTDIR $HOST_FORMATDESTDIR $HOST_DATADEPENDENTDIR
rm -rf $HOST_IMPLICITSEQDESTDIR $HOST_VARIMPLICITSEQDESTDIR $HOST_DYNASEQIMPLICITDESTDIR
rm -rf $HOST_VARINITIALSTATEIMPLICITSEQDESTDIR $HOST_INITIALSTATEIMPLICITSEQDESTDIR
rm -rf $HOST_SCALARMODELSDESTDIR
rm -rf $HOST_IMAGEMODELSDESTDIR
mkdir -p $HOST_SRCDIR
mkdir -p $HOST_DESTDIR
mkdir -p $HOST_VARDESTDIR
mkdir -p $HOST_IDENTITYDESTDIR
mkdir -p $HOST_SIGDEFDESTDIR
mkdir -p $HOST_TFPARAMETERSDESTDIR
mkdir -p $HOST_IDENTITYBIGDESTDIR
mkdir -p $HOST_SHAPEDESTDIR
mkdir -p $HOST_RESHAPEDESTDIR
mkdir -p $HOST_SEQDESTDIR
mkdir -p $HOST_DYNASEQDESTDIR
mkdir -p $HOST_DYNASEQIMPLICITDESTDIR
mkdir -p $HOST_VARSEQDESTDIR
mkdir -p $HOST_ENSEMBLEDESTDIR
mkdir -p $HOST_NOSHAPEDESTDIR
mkdir -p $HOST_PLGDESTDIR
mkdir -p $HOST_RAGGEDDESTDIR
mkdir -p $HOST_FORMATDESTDIR
mkdir -p $HOST_DATADEPENDENTDIR
mkdir -p $HOST_IMPLICITSEQDESTDIR
mkdir -p $HOST_VARIMPLICITSEQDESTDIR
mkdir -p $HOST_INITIALSTATEIMPLICITSEQDESTDIR
mkdir -p $HOST_VARINITIALSTATEIMPLICITSEQDESTDIR
mkdir -p $HOST_TORCHTRTDESTDIR
mkdir -p $HOST_SCALARMODELSDESTDIR
mkdir -p $HOST_IMAGEMODELSDESTDIR

# Since the models required by ensemble models may not be available
# at this point, storing ensemble models separately so that other qa directories
# are proper model repositories
mkdir -p $HOST_ENSEMBLEDESTDIR/qa_model_repository
mkdir -p $HOST_ENSEMBLEDESTDIR/qa_variable_model_repository
mkdir -p $HOST_ENSEMBLEDESTDIR/qa_identity_model_repository
mkdir -p $HOST_ENSEMBLEDESTDIR/qa_reshape_model_repository
mkdir -p $HOST_ENSEMBLEDESTDIR/qa_sequence_model_repository
mkdir -p $HOST_ENSEMBLEDESTDIR/qa_variable_sequence_model_repository

rm -fr $HOST_SRCDIR
mkdir -p $HOST_SRCDIR

cp ./gen_qa_models.py $HOST_SRCDIR/.
cp ./gen_qa_identity_models.py $HOST_SRCDIR/.
cp ./gen_qa_reshape_models.py $HOST_SRCDIR/.
cp ./gen_qa_noshape_models.py $HOST_SRCDIR/.
cp ./gen_qa_sequence_models.py $HOST_SRCDIR/.
cp ./gen_qa_implicit_models.py $HOST_SRCDIR/.
cp ./gen_qa_dyna_sequence_models.py $HOST_SRCDIR/.
cp ./gen_qa_dyna_sequence_implicit_models.py $HOST_SRCDIR/.
cp ./gen_ensemble_model_utils.py $HOST_SRCDIR/.
cp ./gen_qa_trt_plugin_models.py $HOST_SRCDIR/.
cp ./gen_qa_trt_format_models.py $HOST_SRCDIR/.
cp ./gen_qa_trt_data_dependent_shape.py $HOST_SRCDIR/.
cp ./gen_qa_torchtrt_models.py $HOST_SRCDIR/.
cp ./gen_qa_ragged_models.py $HOST_SRCDIR/.
cp ./test_util.py $HOST_SRCDIR/.
cp ./gen_tag_sigdef.py $HOST_SRCDIR/.
cp ./gen_qa_tf_parameters.py $HOST_SRCDIR/.
cp ./gen_common.py $HOST_SRCDIR/.
cp ./gen_qa_ort_scalar_models.py $HOST_SRCDIR/.
cp ./gen_qa_image_models.py $HOST_SRCDIR/.

ONNXSCRIPT=onnx_gen.cmds
OPENVINOSCRIPT=openvino_gen.cmds
TORCHSCRIPT=torch_gen.cmds
TFSCRIPT=tf_gen.cmds
TRTSCRIPT=trt_gen.cmds

SRCDIR=/tmp/src
DESTDIR=/tmp/models
VARDESTDIR=/tmp/varmodels
IDENTITYDESTDIR=/tmp/zeromodels
SIGDEFDESTDIR=/tmp/sigdefmodels
TFPARAMETERSDIR=/tmp/tfparameters
IDENTITYBIGDESTDIR=/tmp/zerobigmodels
SHAPEDESTDIR=/tmp/shapetensormodels
RESHAPEDESTDIR=/tmp/reshapemodels
SEQDESTDIR=/tmp/seqmodels
IMPLICITSEQDESTDIR=/tmp/implicitseqmodels
INITIALSTATEIMPLICITSEQDESTDIR=/tmp/initialstateimplicitseqmodels
VARIMPLICITSEQDESTDIR=/tmp/varimplicitseqmodels
VARINITIALSTATEIMPLICITSEQDESTDIR=/tmp/varinitialstateimplicitseqmodels
DYNASEQDESTDIR=/tmp/dynaseqmodels
DYNASEQIMPLICITDESTDIR=/tmp/dynaseqimplicitmodels
VARSEQDESTDIR=/tmp/varseqmodels
ENSEMBLEDESTDIR=/tmp/ensemblemodels
NOSHAPEDESTDIR=/tmp/noshapemodels
PLGDESTDIR=/tmp/pluginmodels
RAGGEDDESTDIR=/tmp/raggedmodels
FORMATDESTDIR=/tmp/formatmodels
DATADEPENDENTDIR=/tmp/datadependentmodels
TORCHTRTDESTDIR=/tmp/torchtrtmodels
SCALARMODELSDESTDIR=/tmp/scalarmodels
IMAGEMODELSDESTDIR=/tmp/imagemodels

# OPENVINO
#
# OpenVINO is not available on ARM so skip
if [[ "aarch64" != $(uname -m) ]] ; then

cat >$HOST_SRCDIR/$OPENVINOSCRIPT <<EOF
#!/bin/bash
nvidia-smi -L || true
nvidia-smi || true
set -e
set -x
export DEBIAN_FRONTEND=noninteractive
apt-get update && \
    apt-get install -y --no-install-recommends build-essential cmake libprotobuf-dev \
            protobuf-compiler python3 python3-dev python3-pip wget gnupg2 \
            software-properties-common
ln -s /usr/bin/python3 /usr/bin/python

pip3 install  "numpy<=1.23.5" setuptools

pip3 install openvino==$OPENVINO_VERSION

# Since variable shape tensors are not allowed, identity models may fail to generate.
# TODO Add variable size tensor models after DLIS-2827 adds support for variable shape tensors.
# TODO Add sequence models after DLIS-2864 adds support for sequence/control inputs.
python3 $SRCDIR/gen_qa_models.py --openvino --models_dir=$DESTDIR
chmod -R 777 $DESTDIR
# python3 $SRCDIR/gen_qa_identity_models.py --openvino --models_dir=$IDENTITYDESTDIR
# chmod -R 777 $IDENTITYDESTDIR
python3 $SRCDIR/gen_qa_reshape_models.py --openvino --models_dir=$RESHAPEDESTDIR
chmod -R 777 $RESHAPEDESTDIR
# python3 $SRCDIR/gen_qa_sequence_models.py --openvino --models_dir=$SEQDESTDIR
# chmod -R 777 $SEQDESTDIR
# python3 $SRCDIR/gen_qa_dyna_sequence_models.py --openvino --models_dir=$DYNASEQDESTDIR
# chmod -R 777 $DYNASEQDESTDIR
EOF

chmod a+x $HOST_SRCDIR/$OPENVINOSCRIPT
if [ $? -ne 0 ]; then
    echo -e "Failed: chmod"
    exit 1
fi

docker pull $UBUNTU_IMAGE
docker run $DOCKER_GPU_ARGS --rm --entrypoint $SRCDIR/$OPENVINOSCRIPT \
       --mount type=bind,source=$HOST_SRCDIR,target=$SRCDIR \
       --mount type=bind,source=$HOST_DESTDIR,target=$DESTDIR \
       --mount type=bind,source=$HOST_VARDESTDIR,target=$VARDESTDIR \
       --mount type=bind,source=$HOST_IDENTITYDESTDIR,target=$IDENTITYDESTDIR \
       --mount type=bind,source=$HOST_RESHAPEDESTDIR,target=$RESHAPEDESTDIR \
       --mount type=bind,source=$HOST_SEQDESTDIR,target=$SEQDESTDIR \
       --mount type=bind,source=$HOST_DYNASEQDESTDIR,target=$DYNASEQDESTDIR \
       --mount type=bind,source=$HOST_VARSEQDESTDIR,target=$VARSEQDESTDIR \
       --mount type=bind,source=$HOST_RAGGEDDESTDIR,target=$RAGGEDDESTDIR \
       --mount type=bind,source=$HOST_SCALARMODELSDESTDIR,target=$SCALARMODELSDESTDIR \
       $UBUNTU_IMAGE
if [ $? -ne 0 ]; then
    echo -e "Failed"
    exit 1
fi

fi # [[ "aarch64" != $(uname -m) ]]

# ONNX
cat >$HOST_SRCDIR/$ONNXSCRIPT <<EOF
#!/bin/bash
nvidia-smi -L || true
nvidia-smi || true
set -e
set -x
export DEBIAN_FRONTEND=noninteractive
apt-get update && \
        apt-get install -y --no-install-recommends build-essential cmake libprotobuf-dev \
                protobuf-compiler python3 python3-dev python3-pip
ln -s /usr/bin/python3 /usr/bin/python

pip3 install "protobuf<=3.20.1"  "numpy<=1.23.5" # TODO: Remove current line DLIS-3838
pip3 install --upgrade onnx==${ONNX_VERSION}

python3 $SRCDIR/gen_qa_models.py --onnx --onnx_opset=$ONNX_OPSET --models_dir=$DESTDIR
chmod -R 777 $DESTDIR
python3 $SRCDIR/gen_qa_models.py --onnx --onnx_opset=$ONNX_OPSET --variable --models_dir=$VARDESTDIR
chmod -R 777 $VARDESTDIR
python3 $SRCDIR/gen_qa_identity_models.py --onnx --onnx_opset=$ONNX_OPSET --models_dir=$IDENTITYDESTDIR
chmod -R 777 $IDENTITYDESTDIR
python3 $SRCDIR/gen_qa_reshape_models.py --onnx --onnx_opset=$ONNX_OPSET --variable --models_dir=$RESHAPEDESTDIR
chmod -R 777 $RESHAPEDESTDIR
python3 $SRCDIR/gen_qa_sequence_models.py --onnx --onnx_opset=$ONNX_OPSET --models_dir=$SEQDESTDIR
chmod -R 777 $SEQDESTDIR
python3 $SRCDIR/gen_qa_sequence_models.py --onnx --onnx_opset=$ONNX_OPSET --variable --models_dir=$VARSEQDESTDIR
chmod -R 777 $VARSEQDESTDIR
python3 $SRCDIR/gen_qa_implicit_models.py --onnx --initial-state zero --onnx_opset=$ONNX_OPSET --models_dir=$INITIALSTATEIMPLICITSEQDESTDIR
chmod -R 777 $INITIALSTATEIMPLICITSEQDESTDIR
python3 $SRCDIR/gen_qa_implicit_models.py --onnx --initial-state zero --onnx_opset=$ONNX_OPSET --variable --models_dir=$VARINITIALSTATEIMPLICITSEQDESTDIR
chmod -R 777 $VARINITIALSTATEIMPLICITSEQDESTDIR
python3 $SRCDIR/gen_qa_implicit_models.py --onnx --onnx_opset=$ONNX_OPSET --models_dir=$IMPLICITSEQDESTDIR
chmod -R 777 $IMPLICITSEQDESTDIR
python3 $SRCDIR/gen_qa_implicit_models.py --onnx --onnx_opset=$ONNX_OPSET --variable --models_dir=$VARIMPLICITSEQDESTDIR
chmod -R 777 $VARIMPLICITSEQDESTDIR
python3 $SRCDIR/gen_qa_dyna_sequence_models.py --onnx --onnx_opset=$ONNX_OPSET --models_dir=$DYNASEQDESTDIR
chmod -R 777 $DYNASEQDESTDIR
python3 $SRCDIR/gen_qa_dyna_sequence_implicit_models.py --onnx --onnx_opset=$ONNX_OPSET --models_dir=$DYNASEQIMPLICITDESTDIR
chmod -R 777 $DYNASEQIMPLICITDESTDIR
python3 $SRCDIR/gen_qa_ragged_models.py --onnx --onnx_opset=$ONNX_OPSET --models_dir=$RAGGEDDESTDIR
chmod -R 777 $RAGGEDDESTDIR
python3 $SRCDIR/gen_qa_ort_scalar_models.py --onnx_opset=$ONNX_OPSET --models_dir=$SCALARMODELSDESTDIR
chmod -R 777 $RAGGEDDESTDIR
EOF

chmod a+x $HOST_SRCDIR/$ONNXSCRIPT
if [ $? -ne 0 ]; then
    echo -e "Failed: chmod"
    exit 1
fi

docker pull $UBUNTU_IMAGE
docker run $DOCKER_GPU_ARGS --rm --entrypoint $SRCDIR/$ONNXSCRIPT \
       --mount type=bind,source=$HOST_SRCDIR,target=$SRCDIR \
       --mount type=bind,source=$HOST_DESTDIR,target=$DESTDIR \
       --mount type=bind,source=$HOST_VARDESTDIR,target=$VARDESTDIR \
       --mount type=bind,source=$HOST_IDENTITYDESTDIR,target=$IDENTITYDESTDIR \
       --mount type=bind,source=$HOST_RESHAPEDESTDIR,target=$RESHAPEDESTDIR \
       --mount type=bind,source=$HOST_SEQDESTDIR,target=$SEQDESTDIR \
       --mount type=bind,source=$HOST_DYNASEQDESTDIR,target=$DYNASEQDESTDIR \
       --mount type=bind,source=$HOST_DYNASEQIMPLICITDESTDIR,target=$DYNASEQIMPLICITDESTDIR \
       --mount type=bind,source=$HOST_VARSEQDESTDIR,target=$VARSEQDESTDIR \
       --mount type=bind,source=$HOST_RAGGEDDESTDIR,target=$RAGGEDDESTDIR \
       --mount type=bind,source=$HOST_VARIMPLICITSEQDESTDIR,target=$VARIMPLICITSEQDESTDIR \
       --mount type=bind,source=$HOST_IMPLICITSEQDESTDIR,target=$IMPLICITSEQDESTDIR \
       --mount type=bind,source=$HOST_VARINITIALSTATEIMPLICITSEQDESTDIR,target=$VARINITIALSTATEIMPLICITSEQDESTDIR \
       --mount type=bind,source=$HOST_INITIALSTATEIMPLICITSEQDESTDIR,target=$INITIALSTATEIMPLICITSEQDESTDIR \
       --mount type=bind,source=$HOST_SCALARMODELSDESTDIR,target=$SCALARMODELSDESTDIR \
       $UBUNTU_IMAGE
if [ $? -ne 0 ]; then
    echo -e "Failed"
    exit 1
fi

# PyTorch
cat >$HOST_SRCDIR/$TORCHSCRIPT <<EOF
#!/bin/bash
nvidia-smi -L || true
nvidia-smi || true
set -e
set -x
python3 $SRCDIR/gen_qa_models.py --libtorch --models_dir=$DESTDIR
chmod -R 777 $DESTDIR
python3 $SRCDIR/gen_qa_models.py --libtorch --variable --models_dir=$VARDESTDIR
chmod -R 777 $VARDESTDIR
python3 $SRCDIR/gen_qa_identity_models.py --libtorch --models_dir=$IDENTITYDESTDIR
chmod -R 777 $IDENTITYDESTDIR
python3 $SRCDIR/gen_qa_reshape_models.py --libtorch --variable --models_dir=$RESHAPEDESTDIR
chmod -R 777 $RESHAPEDESTDIR
python3 $SRCDIR/gen_qa_sequence_models.py --libtorch --models_dir=$SEQDESTDIR
chmod -R 777 $SEQDESTDIR
python3 $SRCDIR/gen_qa_sequence_models.py --libtorch --variable --models_dir=$VARSEQDESTDIR
chmod -R 777 $VARSEQDESTDIR
python3 $SRCDIR/gen_qa_implicit_models.py --libtorch --models_dir=$IMPLICITSEQDESTDIR
chmod -R 777 $IMPLICITSEQDESTDIR
python3 $SRCDIR/gen_qa_implicit_models.py --libtorch --variable --models_dir=$VARIMPLICITSEQDESTDIR
chmod -R 777 $VARIMPLICITSEQDESTDIR
python3 $SRCDIR/gen_qa_dyna_sequence_models.py --libtorch --models_dir=$DYNASEQDESTDIR
chmod -R 777 $DYNASEQDESTDIR
python3 $SRCDIR/gen_qa_torchtrt_models.py --models_dir=$TORCHTRTDESTDIR
chmod -R 777 $TORCHTRTDESTDIR
python3 $SRCDIR/gen_qa_ragged_models.py --libtorch --models_dir=$RAGGEDDESTDIR
chmod -R 777 $RAGGEDDESTDIR
# Export torchvision image models to ONNX
python3 $SRCDIR/gen_qa_image_models.py --resnet50 --resnet152 --vgg19 --models_dir=$IMAGEMODELSDESTDIR
chmod -R 777 $IMAGEMODELSDESTDIR
EOF

chmod a+x $HOST_SRCDIR/$TORCHSCRIPT
if [ $? -ne 0 ]; then
    echo -e "Failed: chmod"
    exit 1
fi

docker pull $PYTORCH_IMAGE
docker run $DOCKER_GPU_ARGS --rm --entrypoint $SRCDIR/$TORCHSCRIPT \
       --mount type=bind,source=$HOST_SRCDIR,target=$SRCDIR \
       --mount type=bind,source=$HOST_DESTDIR,target=$DESTDIR \
       --mount type=bind,source=$HOST_VARDESTDIR,target=$VARDESTDIR \
       --mount type=bind,source=$HOST_IDENTITYDESTDIR,target=$IDENTITYDESTDIR \
       --mount type=bind,source=$HOST_RESHAPEDESTDIR,target=$RESHAPEDESTDIR \
       --mount type=bind,source=$HOST_SEQDESTDIR,target=$SEQDESTDIR \
       --mount type=bind,source=$HOST_DYNASEQDESTDIR,target=$DYNASEQDESTDIR \
       --mount type=bind,source=$HOST_VARSEQDESTDIR,target=$VARSEQDESTDIR \
       --mount type=bind,source=$HOST_TORCHTRTDESTDIR,target=$TORCHTRTDESTDIR \
       --mount type=bind,source=$HOST_RAGGEDDESTDIR,target=$RAGGEDDESTDIR \
       --mount type=bind,source=$HOST_VARIMPLICITSEQDESTDIR,target=$VARIMPLICITSEQDESTDIR \
       --mount type=bind,source=$HOST_IMPLICITSEQDESTDIR,target=$IMPLICITSEQDESTDIR \
       --mount type=bind,source=$HOST_IMAGEMODELSDESTDIR,target=$IMAGEMODELSDESTDIR \
       $PYTORCH_IMAGE
if [ $? -ne 0 ]; then
    echo -e "Failed"
    exit 1
fi

# Tensorflow
cat >$HOST_SRCDIR/$TFSCRIPT <<EOF
#!/bin/bash
nvidia-smi -L || true
nvidia-smi || true
set -e
set -x
# Segmentation fault with protobuf 4.24.0 (https://github.com/tensorflow/tensorflow/issues/61551)
# Upgrade protobuf version to fix the issue.
pip3 install "protobuf>4.24.0"

python3 $SRCDIR/gen_qa_models.py --graphdef --savedmodel --models_dir=$DESTDIR
chmod -R 777 $DESTDIR
python3 $SRCDIR/gen_qa_models.py --graphdef --savedmodel --variable --models_dir=$VARDESTDIR
chmod -R 777 $VARDESTDIR
python3 $SRCDIR/gen_qa_identity_models.py --graphdef --savedmodel --models_dir=$IDENTITYDESTDIR
chmod -R 777 $IDENTITYDESTDIR
python3 $SRCDIR/gen_qa_reshape_models.py --graphdef --savedmodel --variable --models_dir=$RESHAPEDESTDIR
chmod -R 777 $RESHAPEDESTDIR
python3 $SRCDIR/gen_qa_sequence_models.py --graphdef --savedmodel --models_dir=$SEQDESTDIR
chmod -R 777 $SEQDESTDIR
python3 $SRCDIR/gen_qa_sequence_models.py --graphdef --savedmodel --variable --models_dir=$VARSEQDESTDIR
chmod -R 777 $VARSEQDESTDIR
python3 $SRCDIR/gen_qa_dyna_sequence_models.py --graphdef --savedmodel --models_dir=$DYNASEQDESTDIR
chmod -R 777 $DYNASEQDESTDIR
python3 $SRCDIR/gen_qa_noshape_models.py --savedmodel --models_dir=$NOSHAPEDESTDIR
chmod -R 777 $NOSHAPEDESTDIR
python3 $SRCDIR/gen_qa_ragged_models.py --savedmodel --models_dir=$RAGGEDDESTDIR
chmod -R 777 $RAGGEDDESTDIR
python3 $SRCDIR/gen_qa_models.py --ensemble --models_dir=$ENSEMBLEDESTDIR/qa_model_repository
python3 $SRCDIR/gen_qa_models.py --ensemble --variable --models_dir=$ENSEMBLEDESTDIR/qa_variable_model_repository
python3 $SRCDIR/gen_qa_reshape_models.py --ensemble --models_dir=$ENSEMBLEDESTDIR/qa_reshape_model_repository
python3 $SRCDIR/gen_qa_identity_models.py --ensemble --models_dir=$ENSEMBLEDESTDIR/qa_identity_model_repository
python3 $SRCDIR/gen_qa_sequence_models.py --ensemble --models_dir=$ENSEMBLEDESTDIR/qa_sequence_model_repository
python3 $SRCDIR/gen_qa_sequence_models.py --ensemble --variable --models_dir=$ENSEMBLEDESTDIR/qa_variable_sequence_model_repository
chmod -R 777 $ENSEMBLEDESTDIR
python3 $SRCDIR/gen_tag_sigdef.py --dir $SIGDEFDESTDIR
chmod -R 777 $SIGDEFDESTDIR
python3 $SRCDIR/gen_qa_tf_parameters.py --models_dir $TFPARAMETERSDIR
chmod -R 777 $TFPARAMETERSDIR
EOF

chmod a+x $HOST_SRCDIR/$TFSCRIPT
if [ $? -ne 0 ]; then
    echo -e "Failed: chmod"
    exit 1
fi

docker pull $TENSORFLOW_IMAGE
docker run $DOCKER_GPU_ARGS --rm --entrypoint $SRCDIR/$TFSCRIPT \
       --mount type=bind,source=$HOST_SRCDIR,target=$SRCDIR \
       --mount type=bind,source=$HOST_DESTDIR,target=$DESTDIR \
       --mount type=bind,source=$HOST_VARDESTDIR,target=$VARDESTDIR \
       --mount type=bind,source=$HOST_IDENTITYDESTDIR,target=$IDENTITYDESTDIR \
       --mount type=bind,source=$HOST_SIGDEFDESTDIR,target=$SIGDEFDESTDIR \
       --mount type=bind,source=$HOST_RESHAPEDESTDIR,target=$RESHAPEDESTDIR \
       --mount type=bind,source=$HOST_SEQDESTDIR,target=$SEQDESTDIR \
       --mount type=bind,source=$HOST_DYNASEQDESTDIR,target=$DYNASEQDESTDIR \
       --mount type=bind,source=$HOST_VARSEQDESTDIR,target=$VARSEQDESTDIR \
       --mount type=bind,source=$HOST_NOSHAPEDESTDIR,target=$NOSHAPEDESTDIR \
       --mount type=bind,source=$HOST_ENSEMBLEDESTDIR,target=$ENSEMBLEDESTDIR \
       --mount type=bind,source=$HOST_RAGGEDDESTDIR,target=$RAGGEDDESTDIR \
       --mount type=bind,source=$HOST_TFPARAMETERSDESTDIR,target=$TFPARAMETERSDIR \
       $TENSORFLOW_IMAGE
if [ $? -ne 0 ]; then
    echo -e "Failed"
    exit 1
fi

# TensorRT
cat >$HOST_SRCDIR/$TRTSCRIPT <<EOF
#!/bin/bash
nvidia-smi -L || true
nvidia-smi || true
set -e
set -x
dpkg -l | grep TensorRT
export TRT_SUPPRESS_DEPRECATION_WARNINGS=1
# Models using shape tensor i/o
python3 $SRCDIR/gen_qa_identity_models.py --tensorrt-shape-io --models_dir=$SHAPEDESTDIR
python3 $SRCDIR/gen_qa_sequence_models.py --tensorrt-shape-io --models_dir=$SHAPEDESTDIR
python3 $SRCDIR/gen_qa_dyna_sequence_models.py --tensorrt-shape-io --models_dir=$SHAPEDESTDIR
chmod -R 777 $SHAPEDESTDIR
python3 $SRCDIR/gen_qa_models.py --tensorrt --models_dir=$DESTDIR
chmod -R 777 $DESTDIR
python3 $SRCDIR/gen_qa_models.py --tensorrt --variable --models_dir=$VARDESTDIR
chmod -R 777 $VARDESTDIR
python3 $SRCDIR/gen_qa_identity_models.py --tensorrt --models_dir=$IDENTITYDESTDIR
python3 $SRCDIR/gen_qa_identity_models.py --tensorrt-compat --models_dir=$IDENTITYDESTDIR
chmod -R 777 $IDENTITYDESTDIR
python3 $SRCDIR/gen_qa_identity_models.py --tensorrt-big --models_dir=$IDENTITYBIGDESTDIR
chmod -R 777 $IDENTITYBIGDESTDIR
python3 $SRCDIR/gen_qa_reshape_models.py --tensorrt --variable --models_dir=$RESHAPEDESTDIR
chmod -R 777 $RESHAPEDESTDIR
python3 $SRCDIR/gen_qa_sequence_models.py --tensorrt --models_dir=$SEQDESTDIR
chmod -R 777 $SEQDESTDIR
python3 $SRCDIR/gen_qa_implicit_models.py --tensorrt --models_dir=$IMPLICITSEQDESTDIR
chmod -R 777 $IMPLICITSEQDESTDIR
python3 $SRCDIR/gen_qa_implicit_models.py --tensorrt --variable --models_dir=$VARIMPLICITSEQDESTDIR
chmod -R 777 $VARIMPLICITSEQDESTDIR
python3 $SRCDIR/gen_qa_dyna_sequence_models.py --tensorrt --models_dir=$DYNASEQDESTDIR
chmod -R 777 $DYNASEQDESTDIR
python3 $SRCDIR/gen_qa_sequence_models.py --tensorrt --variable --models_dir=$VARSEQDESTDIR
chmod -R 777 $VARSEQDESTDIR
python3 $SRCDIR/gen_qa_dyna_sequence_implicit_models.py --tensorrt --models_dir=$DYNASEQIMPLICITDESTDIR
chmod -R 777 $DYNASEQIMPLICITDESTDIR
python3 $SRCDIR/gen_qa_ragged_models.py --tensorrt --models_dir=$RAGGEDDESTDIR
chmod -R 777 $RAGGEDDESTDIR
python3 $SRCDIR/gen_qa_trt_format_models.py --models_dir=$FORMATDESTDIR
chmod -R 777 $FORMATDESTDIR
python3 $SRCDIR/gen_qa_trt_data_dependent_shape.py --models_dir=$DATADEPENDENTDIR
chmod -R 777 $DATADEPENDENTDIR
# Make shared library for custom Hardmax plugin.
(git clone -b release/10.0 https://github.com/NVIDIA/TensorRT.git && \
cd /workspace/TensorRT/samples/python/onnx_custom_plugin && rm -rf build && mkdir build && \
cd build && cmake .. && make -j && cp libcustomHardmaxPlugin.so $PLGDESTDIR/.)
LD_PRELOAD=$PLGDESTDIR/libcustomHardmaxPlugin.so python3 $SRCDIR/gen_qa_trt_plugin_models.py --models_dir=$PLGDESTDIR
chmod -R 777 $PLGDESTDIR
EOF

chmod a+x $HOST_SRCDIR/$TRTSCRIPT
if [ $? -ne 0 ]; then
    echo -e "Failed: chmod"
    exit 1
fi

docker pull $TENSORRT_IMAGE
docker run $DOCKER_GPU_ARGS --rm --entrypoint $SRCDIR/$TRTSCRIPT \
       --mount type=bind,source=$HOST_SRCDIR,target=$SRCDIR \
       --mount type=bind,source=$HOST_DESTDIR,target=$DESTDIR \
       --mount type=bind,source=$HOST_VARDESTDIR,target=$VARDESTDIR \
       --mount type=bind,source=$HOST_IDENTITYDESTDIR,target=$IDENTITYDESTDIR \
       --mount type=bind,source=$HOST_IDENTITYBIGDESTDIR,target=$IDENTITYBIGDESTDIR \
       --mount type=bind,source=$HOST_SHAPEDESTDIR,target=$SHAPEDESTDIR \
       --mount type=bind,source=$HOST_RESHAPEDESTDIR,target=$RESHAPEDESTDIR \
       --mount type=bind,source=$HOST_SEQDESTDIR,target=$SEQDESTDIR \
       --mount type=bind,source=$HOST_DYNASEQDESTDIR,target=$DYNASEQDESTDIR \
       --mount type=bind,source=$HOST_VARSEQDESTDIR,target=$VARSEQDESTDIR \
       --mount type=bind,source=$HOST_PLGDESTDIR,target=$PLGDESTDIR \
       --mount type=bind,source=$HOST_RAGGEDDESTDIR,target=$RAGGEDDESTDIR \
       --mount type=bind,source=$HOST_FORMATDESTDIR,target=$FORMATDESTDIR \
       --mount type=bind,source=$HOST_DATADEPENDENTDIR,target=$DATADEPENDENTDIR \
       --mount type=bind,source=$HOST_VARIMPLICITSEQDESTDIR,target=$VARIMPLICITSEQDESTDIR \
       --mount type=bind,source=$HOST_DYNASEQIMPLICITDESTDIR,target=$DYNASEQIMPLICITDESTDIR \
       --mount type=bind,source=$HOST_IMPLICITSEQDESTDIR,target=$IMPLICITSEQDESTDIR \
       $TENSORRT_IMAGE

if [ $? -ne 0 ]; then
    echo -e "Failed"
    exit 1
fi
